import torch
import numpy as np
import pandas as pd 
import matplotlib.pyplot as plt
from sklearn.decomposition import PCA
from sklearn.preprocessing import StandardScaler

# instance_file = "/root/ws/logs/tf-logs/gir_r50_ins/instance_features.pth"
#instance_file = "/root/autodl-tmp/models/girformer_gir_meshb_r50_160k_triloss/instance_features_0004999.pth"
instance_file = "/root/ws/logs/tf-logs/br_stpls3d_r18_160k_tri/instance_features_0159999.pth"


def do_visual_file(file):
            # res = {
            # "mean":result_mean,
            # "std":result_std,
            # "all_features":g_all_feats_tensor,
            # "all_ids":g_all_ids_tensor
            # }
    res = torch.load(file)
    
    ids = torch.unique(res["all_ids"]).numpy()
    
    all_ids = res["all_ids"].numpy()
    all_features = res["all_features"].cpu().numpy()
    
    pca = PCA(n_components=2)
    principalComponents = pca.fit_transform(all_features)
    principalDf = pd.DataFrame(data = principalComponents
             , columns = ['principal component 1', 'principal component 2'])
    
    fig = plt.figure(figsize = (8,8))
    
    ax = fig.add_subplot(1,1,1) 
    # ax.set_xlabel('Principal Component 1', fontsize = 15)
    # ax.set_ylabel('Principal Component 2', fontsize = 15)
    # ax.set_title('2 Component PCA', fontsize = 20)


    colors = np.random.rand(len(ids),3)
    for ind, color in zip(ids,colors):
        indicesToKeep = all_ids == ind
        ax.scatter(principalDf.loc[indicesToKeep, 'principal component 1']
                , principalDf.loc[indicesToKeep, 'principal component 2']
                , c = (color[0],color[1],color[2])
                , s = 5)
    
    plt.axis('off')
    # ax.legend(ids)
    # ax.grid()
    # plt.axes().get_xaxis().set_visible(False) # 隐藏x坐标轴
    # plt.axes().get_yaxis().set_visible(False) # 隐藏y坐标轴
    
    plt.savefig(instance_file + '.png')
    
do_visual_file(instance_file)